#!/usr/bin/perl
use Getopt::Long;
use File::Basename;
use Cwd;
use File::Spec;
use strict;

my $MPIFORMATDB = 1; #use 	MPIFORMATDB to format database
my $NODECOUNT =$ENV{NODECOUNT};

############################################################################
# command line arguments
############################################################################
my $formatdb;

my $seqs;    #number of sequence in input file.
my $input;                #input file
my $nsplit;               #number of split
my $v;            #the number of letter in each segment in million of letters
my $letters;      #the number of letter in the input file
my $p = undef;      #whether it is protein or not.
my $dir;          #the directory where $input it located.


my $bname;        #input base name.
my $title; #title in nal pal file.
my $ooption;
my $diffname;
my $destname;
#main sub routine.
{
	parseOption();
	
	#set $formatdb.
	decideFormatdbCommand();
	chdir $dir or die "can not change to $dir";
	formatdData();
	unlink("mpiblast_temp.conf");
	exit(0);
}

sub parseOption
{
	Getopt::Long::Configure( ( "pass_through", "no_auto_abbrev" ) );
	my $version;
	my $noemboss;
	GetOptions(
				"i=s"      => \$input,
				"split=i"  => \$nsplit,
				"v=i"      => \$v,
				"p=s"      => \$p,
				"X=s"      => \$letters,
				"t=s"=>\$title,
				"o=s"=>\$ooption,
				"n=s"=>\$diffname,
				"dest=s"=>\$destname,
	);
	
	
	usage() if ( !defined($input) );
	usage() if (!defined($p));
	if (!defined($ooption))
	{
		$ooption="F";
	}
	#normalize input file.
	$dir = dirname($input);
	  if ($dir eq "." &&  $^O eq 'MSWin32' )
	  {
	      print STDERR "Please use absolute path for sequence file $input\n";
	      exit;
	  } else
	  {
	      $dir = cwd() if ( $dir eq "." || $dir eq "" );

	  }

	$bname = basename($input);
	unless ( -e $input )
	{
		print STDERR "Input data file $input not found - abort.\n";
		exit(1);
	}	
	if ($destname)
	{
		$diffname=$destname;
		$title=$destname;
	}
}
#################################################################
sub decideFormatdbCommand
{
	$formatdb = "formatdb";
	if ($MPIFORMATDB)
	{
		my $mpiformatdb=`which mpiformatdb`;
		if ($mpiformatdb)
		{
			$formatdb=$mpiformatdb
		} else
		{
			$MPIFORMATDB = 0;
		}
	}
}

sub formatdData
{
	
#maybe used by ncbi as tmp directory. However, I found this setting is not effective in OSX.
#the OS temporary directory may not be enough for formatdb.
	$ENV{TMPDIR} = $dir;
	my $split = undef;
	if ( !defined($nsplit) && !defined($v) )
	{
		if (!$NODECOUNT)
		{
			#default to 8
			$split = 8;
		} else
		{
			$split = $NODECOUNT;
		}
	}
	else
	{
		$split = $nsplit;
	}
	
	#$split = 8 if ( $split > 8 );
	
		if ( !$letters )
		{

			#count letter only if we do not have it from command line
			( $letters, $seqs ) = count_letters($input);
		}
	
			( $v, $letters ) = internal_determine_v( $letters, $split, $v);
		
	my $cmd = "";
	if ($MPIFORMATDB)
	{
		$ENV{MPIBLAST_SHARED} = $dir;

		#disable local cache
		$ENV{MPIBLAST_LOCAL} = $dir;
		my $mpiconfig ="mpiblast_temp.conf";
		#always regenerate this one.
		my $fh;
		open( $fh, ">", $mpiconfig )	  or die "can not open $mpiconfig for writing:$!";
		print $fh "$dir\n";
		print $fh "$dir\n";
		close($fh);


	
		$cmd =
"$formatdb --skip-reorder  --frag-size=$v --quiet -p $p -i $bname  --config-file=$mpiconfig -o $ooption ";
		if ($title)
		{
			$cmd.="-t $title ";
		}
		if ($diffname)
		{
			$cmd.="-n $diffname "
		}
		if ( $v > 1000 )
		{

			#this is a big dataset source file, using sparse option.
			$cmd .= "-s T @ARGV";
		}
		else
		{
			$cmd .= "@ARGV";
		}
	}
	else
	{
		$cmd="$formatdb -v $v -p $p -i $bname -o $ooption ";
		if ($title)
		{
			$cmd.="-t $title ";
		}
		if ($diffname)
		{
			$cmd.="-n $diffname "
		}
			if ( $v > 1000 )
		{

			#this is a big dataset source file, using sparse option.
			$cmd .= "-s T @ARGV";
		}
		else
		{
			$cmd .= "@ARGV";
		}
	}
	print STDERR "Format command:  $cmd\n";

	my $output = `$cmd 2>&1`;
	if ( $? != 0 )
	{
		die $output;
	}
	if ($MPIFORMATDB)
	{
		my @mpffiles = glob("$dir/$bname*mbf");
		if (@mpffiles)
		{
			chmod 0764, @mpffiles;
		}
	}

		
}


# returns the most appropriate v parameter
sub internal_determine_v
{
	my $l      = shift;
	my $pieces = shift;    # number of pieces
	my $vparam = shift;    #the volume passed from command line argument
#	print "$l, $pieces, $vparam\n";
	if ($vparam)
	{
	    if ($vparam>4000)
	    {
		print STDERR "volume is too big. The maximal value of volume is 4000\n";
		exit;
	    }
	    return $vparam;
	}
	my $v = 0;
	$v = $l;
	$v= $l/$pieces;

	#the formatdb unit is in million
	$v = int( $v / (1000000) ) + 1;
	if ($v>4000)
	{
	    my $worksplit=$l/4000000000;
	    $worksplit=int($worksplit)+1;
	    print STDERR "the number of split is too small. I will use $worksplit number of splits\n";
	    $v=4000;
	}

	return ( $v, $l );
}



sub count_letters
{

	# arguments:
	my $ffile  = shift;
	my $zipped = 0;
	if ( scalar(@_) > 0 )
	{
		$zipped = shift;
	}
	my $fh = undef;
	if ($zipped)
	{
		$fh = new IO::Zlib;
		$fh->open( $ffile, "rb" ) or die "can't open $ffile:$!";
	}
	else
	{
		open( $fh, "$ffile" ) or die "can't open $ffile:$!";
	}
	my $totalSeqs    = 0;
	my $totalLetters = 0;
	while (<$fh>)
	{
		if (/^>/)
		{
			$totalSeqs++;
			next;
		}
		else
		{
			$totalLetters += ( length($_) - 1 );
		}
	}
	close $fh;
	return ( $totalLetters, $totalSeqs );
}



sub usage
{
	print "
btformatdb -i <sequence file>
         -p T/F 
	 [-o T/F ]
         [--dest <different database name> ]
         [ --split <# of segments>]
         [ -v volume ]     
         [ other formatdb option ]

	sequence file: sequence file in fasta format
	-p T: the input file is protein sequence
	   F: the input file is nucleotide sequence
    --dest different database name. If you want the resulting database name different from 
		input file name, please use this option.
	-o T/F: If sequences in input file have unique index, please specify the 
		-o T option from formatdb. By default, -o F is assumed.

	For detailed explanation for option, please refer to formatdb document.
	
	

By default, mpiformatdb is used to format the file
	mpiformatdb is used if mpiformatdb binary can not found.
        You can specify the number of segments desired with --split option.
        --split is ignored if -v is specified.
	If --split is not specified, the number of split is default to 
            NODECOUNT taken from ENVIRONMENTAL variables as the number of segments.
	NODECOUNT is default to 8 if not found in ENVIRONMENTAL vaiables
	volume's unit is in million of letters.
        
        The formatted segments will be in  the directory where the input file is located.
";
	exit(1);
}
